from enum import Enum
from typing import Dict, Any, List, Union, Protocol, overload, Type


# Define protocol for hook callbacks
class HookCallback(Protocol):
    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...


class BaseHookEvents(str, Enum):
    """Base class for hook events"""

    pass


class HookRegistry:
    """Type-safe hook registry with configurable event types.

    This class provides a way to register and manage callback functions
    that are triggered at specific points during execution.
    """

    def __init__(self, event_enum: Type[BaseHookEvents]) -> None:
        """Initialize registry with specific event types.

        Args:
            event_enum: Enum class defining the available hook events
        """
        self.event_enum = event_enum
        # Initialize hook storage
        self._hooks: Dict[BaseHookEvents, List[HookCallback]] = {
            event: [] for event in event_enum
        }

    @overload
    def __getitem__(self, key: BaseHookEvents) -> List[HookCallback]: ...

    @overload
    def __getitem__(self, key: str) -> List[HookCallback]: ...

    def __getitem__(self, key: Union[BaseHookEvents, str]) -> List[HookCallback]:
        """Get hooks for the specified event.

        Args:
            key: Event identifier, either as enum value or string

        Returns:
            List of callback functions registered for that event

        Raises:
            KeyError: If the event is not found
        """
        if isinstance(key, str):
            try:
                key = self.event_enum(key)  # Use the configured enum type
            except ValueError:
                raise KeyError(f"Unknown hook event name: {key}")

        return self._hooks[key]

    def __setitem__(
        self, key: Union[BaseHookEvents, str], value: List[HookCallback]
    ) -> None:
        """Set hooks for the specified event.

        Args:
            key: Event identifier, either as enum value or string
            value: List of callback functions to register

        Raises:
            TypeError: If value is not a list
            KeyError: If the event is not found
        """
        if not isinstance(value, list):
            raise TypeError("Hook callbacks must be stored in a list")

        if isinstance(key, str):
            try:
                key = self.event_enum(key)  # Use the configured enum type
            except ValueError:
                raise KeyError(f"Unknown hook event name: {key}")

        self._hooks[key] = value

    def register(
        self, event: Union[BaseHookEvents, str], callback: HookCallback
    ) -> None:
        """Register a callback for a specific event.

        Args:
            event: Event identifier, either as enum value or string
            callback: Function to call when the event occurs

        Raises:
            KeyError: If the event is not found
        """
        self[event].append(callback)

    def execute(
        self, event: Union[BaseHookEvents, str], *args: Any, **kwargs: Any
    ) -> List[Any]:
        """Execute all callbacks for a specific event.

        Args:
            event: Event to trigger
            *args, **kwargs: Arguments to pass to callbacks

        Returns:
            List of results from all callbacks
        """
        return [callback(*args, **kwargs) for callback in self[event]]

    def remove(self, event: Union[BaseHookEvents, str], callback: HookCallback) -> bool:
        """Remove a callback from a specific event.

        Args:
            event: Event identifier
            callback: The callback function to remove

        Returns:
            bool: True if callback was removed, False if not found
        """
        try:
            callbacks = self[event]
            if callback in callbacks:
                callbacks.remove(callback)
                return True
            return False
        except KeyError:
            return False


class TrainHookEvents(BaseHookEvents):
    """Standard training hook events"""

    BEFORE_TRAIN = "before_train"
    AFTER_TRAIN = "after_train"
    BEFORE_TASK = "before_task"
    AFTER_TASK = "after_task"
    BEFORE_EPOCH = "before_epoch"
    AFTER_EPOCH = "after_epoch"
    BEFORE_BATCH = "before_batch"
    AFTER_BATCH = "after_batch"
    BEFORE_FORWARD = "before_forward"
    AFTER_FORWARD = "after_forward"
    BEFORE_BACKWARD = "before_backward"
    AFTER_BACKWARD = "after_backward"
    BEFORE_OPTIMIZER_STEP = "before_optimizer_step"
    AFTER_OPTIMIZER_STEP = "after_optimizer_step"

    # Additional hooks from HookRegistry
    EXTRACT_FEATURES = "extract_features"
    PROCESS_FEATURES = "process_features"
    BEFORE_CLASSIFIER = "before_classifier"
    AFTER_CLASSIFIER = "after_classifier"
    BEFORE_EVAL = "before_eval"
    AFTER_EVAL = "after_eval"
    BEFORE_MEMORY_UPDATE = "before_memory_update"
    AFTER_MEMORY_UPDATE = "after_memory_update"


class TrainHookManager:
    """
    Hook manager for extensibility in the training process.

    This class provides a convenient interface to register and execute hooks at
    different points in the learning process, allowing for extensibility
    without modifying core code.

    Attributes:
        hooks (HookRegistry): Registry for managing callbacks
    """

    def __init__(self):
        self.hooks = HookRegistry(TrainHookEvents)

    def register_hook(
        self, hook_point: Union[TrainHookEvents, str], callback: HookCallback
    ) -> None:
        """
        Register a callback function for a specific hook point.

        Args:
            hook_point: Event identifier, either as TrainHookEvents enum or string
            callback: Function to call when the event occurs

        Raises:
            KeyError: If the hook_point is not found in TrainHookEvents
        """
        self.hooks.register(hook_point, callback)

    def remove_hook(
        self, hook_point: Union[TrainHookEvents, str], callback: HookCallback
    ) -> bool:
        """
        Remove a callback from a specific hook point.

        Args:
            hook_point: Event identifier, either as TrainHookEvents enum or string
            callback: The callback function to remove

        Returns:
            bool: True if callback was removed, False if not found

        Raises:
            KeyError: If the hook_point is not found in TrainHookEvents
        """
        try:
            hook_callbacks = self.hooks[hook_point]
            if callback in hook_callbacks:
                hook_callbacks.remove(callback)
                return True
            return False
        except KeyError:
            return False

    def call_hook(
        self, hook_point: Union[TrainHookEvents, str], *args: Any, **kwargs: Any
    ) -> List[Any]:
        """
        Call all callbacks registered for a hook point.

        Args:
            hook_point: Event to trigger
            *args, **kwargs: Arguments to pass to callbacks

        Returns:
            List of results from all callbacks

        Raises:
            KeyError: If the hook_point is not found in TrainHookEvents
        """
        return self.hooks.execute(hook_point, *args, **kwargs)

    def get_hooks(self, hook_point: Union[TrainHookEvents, str]) -> List[HookCallback]:
        """
        Get all callbacks registered for a specific hook point.

        Args:
            hook_point: Event identifier, either as TrainHookEvents enum or string

        Returns:
            List of callback functions registered for that event

        Raises:
            KeyError: If the hook_point is not found in TrainHookEvents
        """
        return self.hooks[hook_point]

    def decorate_with_hooks(self, method, before_hook, after_hook):
        """
        Utility method to decorate a method with hooks.

        This is an alternative to using the @with_hooks decorator directly.

        Args:
            method: The method to decorate
            before_hook: Hook to call before method execution
            after_hook: Hook to call after method execution

        Returns:
            Decorated method with hooks
        """
        from __temporary__.__backups.hook_decorators import with_hooks
        return with_hooks(before_hook, after_hook)(method)

    def wrap_batch_processing(self, method, hook_name):
        """
        Utility method to wrap a batch processing method with a hook.

        Args:
            method: The batch processing method to wrap
            hook_name: Hook to call during processing

        Returns:
            Wrapped method with batch hook handling
        """
        from __temporary__.__backups.hook_decorators import batch_hook_wrapper
        return batch_hook_wrapper(hook_name)(method)
